# encoding=utf-8
import os
import pickle
import tempfile
import SimpleITK as sitk
import cv2
import numpy as np
import NNUnetTools
from RunHelper import RunHelper
from nnunet.training.model_restore import load_model_and_checkpoint_files
from nnunet.preprocessing.preprocessing import GenericPreprocessor


class NNUnetV1Preprocessor:
    def __init__(self, pkl_path: str):
        with open(pkl_path, "rb") as f:
            self.plans = pickle.load(f)
            d_n = self.plans['normalization_schemes']
            d_m = self.plans['use_mask_for_norm']
            d_t = self.plans['transpose_forward']
            d_i = self.plans['dataset_properties']['intensityproperties']

            self.preprocessor = GenericPreprocessor(d_n, d_m, d_t, d_i)

    def convert_to_nii(self, cv_path: str):
        return NNUnetTools.convert2nifti(cv_path)

    def preprocess(self, cv_path: str):
        stage = list(self.plans['plans_per_stage'].keys())[0]
        stage_plans = self.plans['plans_per_stage'][stage]['current_spacing']

        nii_token_path = self.convert_to_nii(cv_path)
        d, _, properties = self.preprocessor.preprocess_test_case(nii_token_path, stage_plans)
        return d, properties


class NNUnetV1Helper(RunHelper):
    def __init__(self, model_path: str,
                 check_point_name: str = 'model_best',
                 fold_num: int = 4,
                 mixed_precision: bool = True):
        # base info
        super().__init__(model_path, None, None)
        self.model_path = model_path
        self.fold_num = fold_num
        self.check_point_name = check_point_name
        self.mixed_precision = mixed_precision

        # preprocessor
        self.preprocessor = NNUnetV1Preprocessor(
            os.path.join(self.model_path, "plans.pkl")
        )

        # runner
        self.trainer, self.params = load_model_and_checkpoint_files(
            self.model_path, self.fold_num,
            mixed_precision=self.mixed_precision,
            checkpoint_name=self.check_point_name
        )
        self.trainer.load_checkpoint_ram(self.params[0], False)
        self.trainer.network.eval()

    def infer(self, data_in):
        # 数据产生
        preprocess_data, properties = self.preprocessor.preprocess(data_in)

        # 生成
        out_data = self.trainer.predict_preprocessed_data_return_seg_and_softmax(
            preprocess_data, do_mirroring=True, mirror_axes=self.trainer.data_aug_params['mirror_axes'],
            use_sliding_window=True, step_size=0.5, use_gaussian=True, all_in_gpu=False,
            mixed_precision=True
        )

        # 转换
        transpose_forward = self.trainer.plans.get('transpose_forward')
        if transpose_forward is not None:
            transpose_backward = self.trainer.plans.get('transpose_backward')
            softmax = out_data[1].transpose([0] + [i + 1 for i in transpose_backward])
        else:
            softmax = out_data[1]

        tmp_path = tempfile.NamedTemporaryFile(suffix='.nii.gz')
        each_obj = softmax[0: 2]  # 0和1
        seg_old_spacing = each_obj.argmax(0)
        seg_resized_itk = sitk.GetImageFromArray(seg_old_spacing.astype(np.uint8))
        seg_resized_itk.SetSpacing(properties['itk_spacing'])
        seg_resized_itk.SetOrigin(properties['itk_origin'])
        seg_resized_itk.SetDirection(properties['itk_direction'])
        sitk.WriteImage(seg_resized_itk, tmp_path.name)

        convert_mat = NNUnetTools.convert2png(tmp_path.name)

        # 恢复相对位置
        ori_mat = cv2.imread(data_in)
        copy_mask = np.zeros(ori_mat.shape[:2], np.uint8)
        h_s, h_e = properties['crop_bbox'][1]
        w_s, w_e = properties['crop_bbox'][2]
        copy_mask[h_s: h_e, w_s: w_e] = convert_mat
        return convert_mat * 255


class NNUnetV1MultiHelper(NNUnetV1Helper):
    def __init__(self,
                 model_path: str,
                 check_point_name: str = 'model_best',
                 fold_num: int = 4,
                 mixed_precision: bool = True):
        super().__init__(model_path, check_point_name, fold_num, mixed_precision)
        self.plans_path = os.path.join(model_path, "plans.pkl")

        with open(self.plans_path, "rb") as f:
            plans = pickle.load(f)
            self.class_list = plans['all_classes']

    def infer(self, data_in):
        # 数据产生
        preprocess_data, properties = self.preprocessor.preprocess(data_in)

        # 生成
        out_data = self.trainer.predict_preprocessed_data_return_seg_and_softmax(
            preprocess_data, do_mirroring=True, mirror_axes=self.trainer.data_aug_params['mirror_axes'],
            use_sliding_window=True, step_size=0.5, use_gaussian=True, all_in_gpu=False,
            mixed_precision=True
        )

        # 转换
        transpose_forward = self.trainer.plans.get('transpose_forward')
        if transpose_forward is not None:
            transpose_backward = self.trainer.plans.get('transpose_backward')
            softmax = out_data[1].transpose([0] + [i + 1 for i in transpose_backward])
        else:
            softmax = out_data[1]

        tmp_path = tempfile.NamedTemporaryFile(suffix='.nii.gz')

        seg_mask = np.zeros(softmax.shape[1:], dtype=np.uint8)
        for i, c in enumerate(self.class_list):
            seg_mask[softmax[i + 1] > 0.5] = c
        seg_resized_itk = sitk.GetImageFromArray(seg_mask)
        seg_resized_itk.SetSpacing(properties['itk_spacing'])
        seg_resized_itk.SetOrigin(properties['itk_origin'])
        seg_resized_itk.SetDirection(properties['itk_direction'])
        sitk.WriteImage(seg_resized_itk, tmp_path.name)

        convert_mat = NNUnetTools.convert2png(tmp_path.name)

        # 恢复相对位置
        ori_mat = cv2.imread(data_in)
        copy_mask = np.zeros(ori_mat.shape[:2], np.uint8)
        h_s, h_e = properties['crop_bbox'][1]
        w_s, w_e = properties['crop_bbox'][2]
        copy_mask[h_s: h_e, w_s: w_e] = convert_mat
        return convert_mat
